import gradio as gr
import pandas as pd
import numpy as np
import plotly.graph_objects as go
import plotly.express as px
import os

import yaml


def pie_chart(labels, values, title):
    # 数据
    # labels = ['苹果', '香蕉', '橙子', '葡萄']
    # values = [25, 30, 20, 25]
    # 绘制饼图
    fig = px.pie(values=values, names=labels, title=title)
    return fig

def get_yearly_book1_output_data(cfg):
    file_path = f"{os.getcwd()}\\raw_data\\年度公共预算支出及分项.xlsx"
    data = pd.read_excel(file_path, index_col=0, skiprows=3)
    data = data.iloc[6:,:]
    data.columns = [cfg['yearly_long_name_to_short'][s] for s in data.columns]
    return data

def get_monthly_data(cfg):
    dir_path = os.getcwd()
    file_path_input = f"{dir_path}\\raw_data\\monthly_data.xlsx"
    data = pd.read_excel(file_path_input, index_col=0,skiprows=1)
    data = data.iloc[6:,:]
    for s in data.columns:
        print(s)
    data.columns = [cfg['long_name_to_short'][s] for s in data.columns]
    return data

def get_config():
    print(os.getcwd())
    dir_path = os.getcwd()
    file_path = '\\raw_data\\fiscal_config.yaml'  # 替换为你的 YAML 文件路径
    path = f"{dir_path}{file_path}"
    print(path)
    with open(path, 'r', encoding='utf-8') as file:
        data = yaml.safe_load(file) # 使用 safe_load 来加载 YAML 文件内容
    print(data)
    return data


def prepare_yearly_data(data, name):
    df = data[name]
    df = df.dropna()
    df = pd.DataFrame(df)
    df.columns = ['累计值']
    df.index = pd.to_datetime(df.index)
    df['year'] = df.index.year if isinstance(df.index, pd.DatetimeIndex) else [pd.Timestamp(x).year for x in df.index]
    df['month'] = df.index.month if isinstance(df.index, pd.DatetimeIndex) else [pd.Timestamp(x).month for x in df.index]
    df.sort_index(inplace=True)
    return df


def prepare_monthly_data(data, name):
    # name = '一般公共预算收入'
    df = data[name]
    df = df.dropna()
    df = pd.DataFrame(df)
    df.columns = ['累计值']
    df.index = pd.to_datetime(df.index)
    date_range = pd.date_range(start=df.index.min(), end=df.index.max(), freq='M')
    df = df.reindex(date_range)
    df['year'] = df.index.year  # 提取年份
    df['month'] = df.index.month  # 提取月份
    # 按年份分组计算当月值
    df.sort_index(inplace=True)
    df['当月值'] = df.groupby('year')['累计值'].diff()

    for row, se in df.iterrows():
        if se['month'] in [1,2]:
            # v = df[(df['year'] == se['year']) & (df['month'] == 2)]['累计值'].values[0]
            v = df[(df['year'] == se['year']) & (df['month'] == 2)]['累计值'][0]
            df.loc[row, '当月值'] = v / 2
    df = df.replace({None: np.nan})
    df['当月值'] = df['当月值'].round(1)
    df = df.loc[:,['year', 'month', '当月值', '累计值']]
    return df


def plot_data(df, name, year_start, data_type='当月值'): 
    df = df[df['year'] >= year_start]
    x = df.index.to_list()
    y = df[data_type].to_list()
    if  '同比' in data_type or '占比' in name:
        h_hover_text = [f'{c.year}年{c.month}月:{v:.1f}%' for c, v in zip(x, y)]
    else:
        h_hover_text = [f'{c.year}年{c.month}月:{v:.1f}亿元' for c, v in zip(x, y)]

    fig = go.Figure(data=[go.Bar(x=x, y=y,hovertext=h_hover_text)])
    fig.update_layout(title=f"{name}",xaxis_title='类别',yaxis_title='数值')
    df['指标'] = name
    return fig, df

def plot_seasonal_data(df, name, year_start, data_type='当月值'): 
    fig = go.Figure()
    df = df[df['year'] >= year_start]
    for year, df_year in df.groupby('year'):
        x = [f"{c.month}月" for c in df_year.index.to_list()]
        y = df_year[data_type].to_list()
        if '同比' in data_type:
            h_hover_text = [f'{c.year}年{c.month}月:{v:.1f}%' for c, v in zip(df_year.index.to_list(), y)]
        else:
            h_hover_text = [f'{c.year}年{c.month}月:{v:.1f}亿元' for c, v in zip(df_year.index.to_list(), y)]
        fig.add_trace(go.Scatter(x=x, y=y, mode='lines', name=year, hovertext=h_hover_text))
    # 更新布局
    fig.update_layout(
        title=name,
        xaxis_title="日期",
        yaxis_title="值",
        legend_title="图例"
    )
    df['指标'] = name
    return fig, df


class FiscalBook:
    def __init__(self):
        self.cfg = get_config()
        self.data = get_monthly_data(self.cfg)
        self.yearly_out_data = get_yearly_book1_output_data(self.cfg)
        self.first_monthly_input_items = self.cfg['input_item']
        self.first_monthly_output_items = self.cfg['output_item']
        
        self.first_yearly_input_items = self.yearly_out_data.columns.to_list()
        self.first_monthly_items = self.first_monthly_input_items + self.first_monthly_output_items
        self.yearly_items = self.first_yearly_input_items
        
        self.book_types = ['一般公共预算', '政府性基金预算', '国有资本经营预算','社会保险基金预算']
        self.plot_types = ['季节', '趋势']
        self.data_types = ['累计同比', '当月同比', '累计值', '当月值']
        self.data_types_map = {'累计同比': '累计值', '当月同比': '当月值'}
        
        self.input_data_dic = {k:prepare_monthly_data(self.data, k) for k in self.first_monthly_input_items}
        self.output_data_dic = {k:prepare_monthly_data(self.data, k) for k in self.first_monthly_output_items}
        self.yearly_output_data_dic = {k:prepare_yearly_data(self.yearly_out_data, k) for k in self.first_yearly_input_items}
        self.data_dic = self.get_data_dic()
    
    def get_data_dic(self):
        data_dic = {}
        for _, items in self.cfg['monthly_series_analyse_items'].items():
            for item in items:
                df = prepare_monthly_data(self.data, item)
                data_dic[item] = df
        return data_dic
    
    @staticmethod
    def process_item_data(df, item, item_df, year, month=12, input=True):
        item_df = item_df[item_df['year'] >= year]
        item_df = item_df[item_df['month'] == month]
        item_df['date'] = [f"{c.year}年-{c.month}月" for c in item_df.index]
        item_df['item'] = item
        item_df = item_df.loc[:, ['date', '累计值', 'item']]
        item_df.columns = ['date', 'value', 'item']
        if not input:
            item_df['value'] = -item_df['value']
            item_df['io'] = '支出'
        else:
            item_df['io'] = '收入'
        df = pd.concat([df, item_df])
        return df
        
    def get_item_value_data(self, year, month):
        # 整理年度数据,收入和支出
        # date, value, item
        df = pd.DataFrame(columns=['date', 'value', 'item', 'io'])
        for item in self.cfg['input_item']:
            item_df = self.input_data_dic[item]
            df = self.process_item_data(df, item, item_df, year, month=12, input=True)
        for item in self.yearly_items[1:]:
            item_df = self.yearly_output_data_dic[item]
            df = self.process_item_data(df, item, item_df, year, month=12, input=False)
        return df
    
    def plot(self, name, year_start, plot_type = '季节性', data_type='当月值'):
        assert plot_type in self.plot_types
        df = self.data_dic[name]
        if data_type in ['累计同比', '当月同比']:
            df = df[df['year'] >= year_start - 1]
            df[data_type] = None
            col = self.data_types_map[data_type]
            for row, se in df.iloc[12:,:].iterrows():
                year = se['year']
                month = se['month']
                v_last_year = df[(df['year'] == year - 1) & (df['month'] == month)][col].values[0]
                v_this_year = df[(df['year'] == year) & (df['month'] == month)][col].values[0]
                df.loc[row, data_type] = (v_this_year - v_last_year) / v_last_year * 100
            df = df.replace({None: np.nan})
            df[data_type] = df[data_type].round(1)
        if plot_type == '季节':
            return plot_seasonal_data(df, name, year_start, data_type)
        else:
            return plot_data(df, name, year_start, data_type)
        

class FiscalData:
    def __init__(self):
        self.cfg = get_config()
        self.one_two_book_items = list(self.cfg['monthly_series_analyse_items'].keys())
        self.first_book = FiscalBook()
    
    def change_book(self, book_type):
        ss = self.cfg['monthly_series_analyse_items'][book_type]
        r_dropdown = gr.Dropdown(choices=ss,label='行名',value=ss[0],interactive=True)
        return r_dropdown

    @staticmethod
    def plot_pyramid_df(df, y, x,color):
        # y = 'date', x = 'value', color = 'item'
        # y = 'date', x = 'value', color = 'io'
        fig = px.bar(df, y=y, x=x, orientation='h', color=color, barmode='relative', )
        fig.update_layout(
            title=f'收支结构分析{color}',
            xaxis_title='左侧支出，右侧收入',
            yaxis_title='日期'
        )
        return fig
    
    def plot_pyramid(self, year):
        # 获取指定年份的12月数据
        df = self.first_book.get_item_value_data(year, 12)
        # 按日期和io分组，求value的和，并重置索引
        io_df = df.groupby(['date', 'io'])['value'].sum().reset_index()
        # 绘制以日期和value为x轴，item为y轴的柱状图
        fig1 = self.plot_pyramid_df(df, 'date', 'value', 'item')
        # 绘制以日期和value为x轴，io为y轴的柱状图
        fig2 = self.plot_pyramid_df(io_df, 'date', 'value', 'io')
        # 返回两个柱状图
        df_year = df[df['date'] == f"{year}年-12月"]
        df_input = df_year[df_year['io'] == '收入']
        df_output = df_year[df_year['io'] == '支出']
        pie_input = pie_chart(
            pd.Series(df_input['item']).to_list(), 
            pd.Series(df_input['value']).to_list(), 
            f'收入分析{year}年'
            )
        df_output['value'] = -df_output['value']
        pie_output = pie_chart(
            pd.Series(df_output['item']).to_list(), 
            pd.Series(df_output['value']).to_list(), 
            f'支出分析{year}年')
        return fig1, fig2, pie_input, pie_output
    
    def html(self):
        with gr.Blocks() as demo:
            gr.Markdown("# 财政数据分析")
            with gr.Tab("月度数据序列分析"):
                with gr.Row():
                    book_type = gr.Radio(
                        label='账本类型', choices=self.one_two_book_items, value=self.one_two_book_items[0], 
                        scale=2)
                    name = gr.Dropdown(label='指标名称', choices=self.first_book.first_monthly_items, value=self.first_book.first_monthly_items[0],scale=2)
                    year_start = gr.Slider(label='起始年份', minimum=2010, maximum=2023, step=1, value=2020, scale=1)
                with gr.Row():    
                    plot_type = gr.Radio(label='图表类型', choices=self.first_book.plot_types,value=self.first_book.plot_types[0],scale=1)
                    data_type = gr.Radio(label='数据类型', choices=self.first_book.data_types,value=self.first_book.data_types[0],scale=3)
                    plot_button = gr.Button("生成图表", scale=1)
                with gr.Tab("图"):
                    plot = gr.Plot()
                with gr.Tab("表"):
                    df = gr.DataFrame()
                plot_button.click(self.first_book.plot, [name, year_start, plot_type, data_type], [plot, df])
                book_type.change(self.change_book, [book_type], [name])

            with gr.Tab("一般公共预算收支结构分析"):
                gr.Markdown("## 综合分析")
                with gr.Row():
                    year_ana = gr.Slider(label='年份', minimum=2010, maximum=2023, step=1, value=2020, scale=1)
                    pr_button = gr.Button("生成图表", scale=1)


                pr_plot = gr.Plot() # 收支结构
                io_plot = gr.Plot() # 收支总项
                with gr.Row():
                    pie_1 = gr.Plot() # 收入结构
                    pie_2 = gr.Plot() # 支出结构
                
                pr_button.click(self.plot_pyramid, [year_ana], [pr_plot, io_plot, pie_1, pie_2])
        
        return demo


fiscal_data = FiscalData()

fiscal_data.html().launch()